﻿// --------------------------------------------------------------------------------------------------------------------
// <copyright file="MultiMap.cs" company="">
//   
// </copyright>
// <summary>
//   A data structure that contains multiple values for each key.
// </summary>
// --------------------------------------------------------------------------------------------------------------------
namespace Skymate.Collections
{
    using System;
    using System.Collections;
    using System.Collections.Concurrent;
    using System.Collections.Generic;
    using System.Linq;
    using System.Linq.Expressions;

    using Skymate.Extensions;

    /// <summary>
    /// A data structure that contains multiple values for each key.
    /// </summary>
    /// <typeparam name="TKey">
    /// The type of key.
    /// </typeparam>
    /// <typeparam name="TValue">
    /// The type of value.
    /// </typeparam>
    public class Multimap<TKey, TValue> : IEnumerable<KeyValuePair<TKey, ICollection<TValue>>>
    {
        /// <summary>
        ///     The _is readonly.
        /// </summary>
        private readonly bool _isReadonly;

        /// <summary>
        ///     The _items.
        /// </summary>
        private readonly IDictionary<TKey, ICollection<TValue>> _items;

        /// <summary>
        ///     The _list creator.
        /// </summary>
        private readonly Func<ICollection<TValue>> _listCreator;

        /// <summary>
        ///     Initializes a new instance of the <see cref="Multimap{TKey,TValue}" /> class.
        /// </summary>
        public Multimap()
            : this(false)
        {
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="Multimap{TKey,TValue}"/> class.
        /// </summary>
        /// <param name="threadSafe">
        /// The thread safe.
        /// </param>
        internal Multimap(bool threadSafe)
        {
            if (threadSafe)
            {
                this._items = new ConcurrentDictionary<TKey, ICollection<TValue>>();
                this._listCreator = () => new SynchronizedCollection<TValue>();
            }
            else
            {
                this._items = new Dictionary<TKey, ICollection<TValue>>();
                this._listCreator = () => new List<TValue>();
            }
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="Multimap{TKey,TValue}"/> class.
        /// </summary>
        /// <param name="listCreator">
        /// The list creator.
        /// </param>
        public Multimap(Func<ICollection<TValue>> listCreator)
            : this(new Dictionary<TKey, ICollection<TValue>>(), listCreator)
        {
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="Multimap{TKey,TValue}"/> class.
        /// </summary>
        /// <param name="dictionary">
        /// The dictionary.
        /// </param>
        /// <param name="listCreator">
        /// The list creator.
        /// </param>
        internal Multimap(IDictionary<TKey, ICollection<TValue>> dictionary, Func<ICollection<TValue>> listCreator)
        {
            this._items = dictionary;
            this._listCreator = listCreator;
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="Multimap{TKey,TValue}"/> class.
        /// </summary>
        /// <param name="dictionary">
        /// The dictionary.
        /// </param>
        /// <param name="isReadonly">
        /// The is readonly.
        /// </param>
        protected Multimap(IDictionary<TKey, ICollection<TValue>> dictionary, bool isReadonly)
        {
            Guard.ArgumentNotNull(() => dictionary);

            this._items = dictionary;

            if (isReadonly && dictionary != null)
            {
                foreach (var kvp in dictionary)
                {
                    dictionary[kvp.Key] = kvp.Value.AsReadOnly();
                }
            }

            this._isReadonly = isReadonly;
        }

        /// <summary>
        ///     Gets the count of groups/keys.
        /// </summary>
        public int Count
        {
            get
            {
                return this._items.Keys.Count;
            }
        }

        /// <summary>
        ///     Gets the total count of items in all groups.
        /// </summary>
        public int TotalValueCount
        {
            get
            {
                return this._items.Values.Sum(x => x.Count);
            }
        }

        /// <summary>
        /// Gets the collection of values stored under the specified key.
        /// </summary>
        /// <param name="key">
        /// The key.
        /// </param>
        /// <returns>
        /// The <see cref="ICollection"/>.
        /// </returns>
        public virtual ICollection<TValue> this[TKey key]
        {
            get
            {
                if (!this._items.ContainsKey(key))
                {
                    if (!this._isReadonly)
                    {
                        this._items[key] = this._listCreator();
                    }
                    else
                    {
                        return null;
                    }
                }

                return this._items[key];
            }
        }

        /// <summary>
        ///     Gets the collection of keys.
        /// </summary>
        public virtual ICollection<TKey> Keys
        {
            get
            {
                return this._items.Keys;
            }
        }

        /// <summary>
        ///     Gets the collection of collections of values.
        /// </summary>
        public virtual ICollection<ICollection<TValue>> Values
        {
            get
            {
                return this._items.Values;
            }
        }

        /// <summary>
        ///     Returns an enumerator that iterates through the multimap.
        /// </summary>
        /// <returns>An <see cref="IEnumerator" /> object that can be used to iterate through the multimap.</returns>
        IEnumerator IEnumerable.GetEnumerator()
        {
            return this.GetEnumerator();
        }

        /// <summary>
        ///     Returns an enumerator that iterates through the multimap.
        /// </summary>
        /// <returns>An <see cref="IEnumerator" /> object that can be used to iterate through the multimap.</returns>
        public virtual IEnumerator<KeyValuePair<TKey, ICollection<TValue>>> GetEnumerator()
        {
            foreach (var pair in this._items)
            {
                yield return pair;
            }
        }

        /// <summary>
        /// The find.
        /// </summary>
        /// <param name="key">
        /// The key.
        /// </param>
        /// <param name="predicate">
        /// The predicate.
        /// </param>
        /// <returns>
        /// The <see cref="IEnumerable"/>.
        /// </returns>
        public IEnumerable<TValue> Find(TKey key, Expression<Func<TValue, bool>> predicate)
        {
            Guard.ArgumentNotNull(() => key);
            Guard.ArgumentNotNull(() => predicate);

            if (this._items.ContainsKey(key))
            {
                return this._items[key].Where(predicate.Compile());
            }

            return Enumerable.Empty<TValue>();
        }

        /// <summary>
        /// Adds the specified value for the specified key.
        /// </summary>
        /// <param name="key">
        /// The key.
        /// </param>
        /// <param name="value">
        /// The value.
        /// </param>
        public virtual void Add(TKey key, TValue value)
        {
            this.CheckNotReadonly();

            this[key].Add(value);
        }

        /// <summary>
        /// Adds the specified values to the specified key.
        /// </summary>
        /// <param name="key">
        /// The key.
        /// </param>
        /// <param name="values">
        /// The values.
        /// </param>
        public virtual void AddRange(TKey key, IEnumerable<TValue> values)
        {
            this.CheckNotReadonly();

            this[key].AddRange(values);
        }

        /// <summary>
        /// Removes the specified value for the specified key.
        /// </summary>
        /// <param name="key">
        /// The key.
        /// </param>
        /// <param name="value">
        /// The value.
        /// </param>
        /// <returns>
        /// <c>True</c> if such a value existed and was removed; otherwise <c>false</c>.
        /// </returns>
        public virtual bool Remove(TKey key, TValue value)
        {
            this.CheckNotReadonly();

            if (!this._items.ContainsKey(key))
            {
                return false;
            }

            var result = this._items[key].Remove(value);
            if (this._items[key].Count == 0)
            {
                this._items.Remove(key);
            }

            return result;
        }

        /// <summary>
        /// Removes all values for the specified key.
        /// </summary>
        /// <param name="key">
        /// The key.
        /// </param>
        /// <returns>
        /// <c>True</c> if any such values existed; otherwise <c>false</c>.
        /// </returns>
        public virtual bool RemoveAll(TKey key)
        {
            this.CheckNotReadonly();
            return this._items.Remove(key);
        }

        /// <summary>
        ///     Removes all values.
        /// </summary>
        public virtual void Clear()
        {
            this.CheckNotReadonly();
            this._items.Clear();
        }

        /// <summary>
        /// Determines whether the multimap contains any values for the specified key.
        /// </summary>
        /// <param name="key">
        /// The key.
        /// </param>
        /// <returns>
        /// <c>True</c> if the multimap has one or more values for the specified key, otherwise <c>false</c>.
        /// </returns>
        public virtual bool ContainsKey(TKey key)
        {
            return this._items.ContainsKey(key);
        }

        /// <summary>
        /// Determines whether the multimap contains the specified value for the specified key.
        /// </summary>
        /// <param name="key">
        /// The key.
        /// </param>
        /// <param name="value">
        /// The value.
        /// </param>
        /// <returns>
        /// <c>True</c> if the multimap contains such a value; otherwise, <c>false</c>.
        /// </returns>
        public virtual bool ContainsValue(TKey key, TValue value)
        {
            return this._items.ContainsKey(key) && this._items[key].Contains(value);
        }

        /// <summary>
        ///     The check not readonly.
        /// </summary>
        /// <exception cref="NotSupportedException">
        /// </exception>
        private void CheckNotReadonly()
        {
            if (this._isReadonly)
            {
                throw new NotSupportedException("Multimap is read-only.");
            }
        }

        #region Static members

        /// <summary>
        ///     The thread safe.
        /// </summary>
        /// <returns>
        ///     The <see cref="Multimap" />.
        /// </returns>
        public static Multimap<TKey, TValue> ThreadSafe()
        {
            return new Multimap<TKey, TValue>(true);
        }

        /// <summary>
        /// The create from lookup.
        /// </summary>
        /// <param name="source">
        /// The source.
        /// </param>
        /// <returns>
        /// The <see cref="Multimap"/>.
        /// </returns>
        public static Multimap<TKey, TValue> CreateFromLookup(ILookup<TKey, TValue> source)
        {
            Guard.ArgumentNotNull(() => source);

            var map = new Multimap<TKey, TValue>();

            foreach (var group in source)
            {
                map.AddRange(group.Key, group);
            }

            return map;
        }

        #endregion
    }
}